package gen

import (
	"io"
	"math"
	"math/rand"
	"os"
    "cs466/findmotif/motif"
)

// describes one benchmark-generation task as defined in the assignment
type Task struct {
	ICPC        float64
	MotifLength uint
	SeqLength   uint
	SeqCount    uint

	p         motif.PWM
	sequences [][]byte
	offsets   []int
}

// generate n random bytes uniformly by reading from /dev/urandom
func randBytes(n int64) (ret []uint8, err error) {
	f, err := os.Open("/dev/urandom")
	if err != nil {
		return nil, err
	}

	ret = make([]byte, n)
	_, err = io.ReadFull(f, ret)
	if err != nil {
		return nil, err
	}

	return
}

// generate n random nucleotides uniformly based on randBytes
func randNucleotides(n int64) (ret []byte, err error) {
	bytes, err := randBytes(n / 4) // 2 bits of entropy per nucleotide
	if err != nil {
		return nil, err
	}

	ret = make([]byte, 0, n)
	for _, b := range bytes {
		for i := 0; i < 4; b, i = b>>2, i+1 {
			ret = append(ret, motif.BitsToNucleotide[b&0x3])
		}
	}

	return
}

// generate a random PWM with ICPC I, n symbols per position, and length m
func randPWM(I float64, n int64, m int64) (ret motif.PWM) {
	if n != 4 {
		panic("n must be 4 until RandColumn can handle other values")
	}
	target_ic := I * float64(m)
	cumulative_ic := 0.
	ret = make(motif.PWM, m)
	for i := range ret {
		// theoretical upper bound on IC of the rest of the PWM (not including
		// the ith column)
		max_rest := (float64(m) - float64(i) - 1) * math.Log2(float64(n))

		// IC needed for the rest of the matrix (including ith column)
		needed_ic := target_ic - cumulative_ic

		// minimum IC for ith column is their difference (but 0 is the lowest
		// possible)
		min := math.Max(needed_ic-max_rest, 0)

		// maximum IC for ith column is simply needed_ic (but ln(n) is the
		// highest possible)
		max := math.Min(needed_ic, math.Log2(float64(n)))

		// generate a random float in (0,1) and scale it to this range for this
		// column's ICPC
		icpc := min + (max-min)*rand.Float64()

		// and generate a column with that ICPC
		ret[i] = motif.Column(RandColumn(icpc))

		cumulative_ic += icpc
	}

	return
}

func Sample(c motif.Column) byte {
	p := rand.Float64()
	var i int
	for i = 0; p >= c[i]; i = i + 1 {
		p -= c[i]
	}
	return motif.BitsToNucleotide[i]
}

func SampleFull(p motif.PWM, dst []byte) {
	for i := range dst {
		dst[i] = SampleAt(p,i)
	}
}

func SampleAt(p motif.PWM, i int) byte {
	return Sample(p[i])
}

func (t *Task) sampleSequence() (ret []byte, offset int, err error) {
	// generate a random sequence
	seq, err := randNucleotides(int64(t.SeqLength))
	if err != nil {
		return nil, 0, err
	}

	// random starting position for motif that leaves enough room for the full
	// motif length
	offset = int(rand.Int31n(int32(t.SeqLength - t.MotifLength)))

	// overwrite the chosen part of the sequence
	SampleFull(t.p, seq[offset : offset+int(t.MotifLength)])

	return seq, offset, err
}

func (t *Task) sampleSequences() ([][]byte, []int, error) {
	t.sequences = make([][]byte, t.SeqCount)
	t.offsets = make([]int, t.SeqCount)
	var err error
	for i := range t.sequences {
		t.sequences[i], t.offsets[i], err = t.sampleSequence()
		if err != nil {
			return nil, nil, err
		}
	}

	return t.sequences, t.offsets, nil
}

func (t *Task) samplePWM() motif.PWM {
	t.p = randPWM(t.ICPC, 4, int64(t.MotifLength))
	return t.p
}

func (t *Task) Generate() (seqs [][]byte, sites []int, p motif.PWM, err error) {
	p = t.samplePWM()
	seqs, sites, err = t.sampleSequences()
	return
}
